#ifndef __ATOM_TRACE_PHOTON_MAP_H_
#define __ATOM_TRACE_PHOTON_MAP_H_

#include "Color.h"
#include "Vector3.h"
#include "atUtil.h"
#include "Scene.h"
#include "Ray.h"
#include "AABBox.h"

namespace AtomTrace
{
	typedef struct Photon
	{
		Color flux;
		Ray pRay;

		Photon()
		{
			flux.Set(0.0f, 0.0f, 0.0f);
			pRay._o.Set(0.0f, 0.0f, 0.0f);
			pRay._dir.SetUnitY();
		}

		const Photon & operator=(const Photon & iPht)
		{
			flux = iPht.flux;
			pRay = iPht.pRay;
			return *this;
		}

		void Set(const Ray & r)
		{ pRay = r; }

		void SaveToMemory(int addr);
		void LoadFromMemory(int addr);
		void LoadFromMemoryGivenIndex(int index);
        void DebugPrint() const;
	}Photon;
    
	int GetMaxPhotonCapacity();
    int GetMaxNumberOfPhotons();
    
    int GetSingleEmittedPhotons();
    
	Photon LoadPhotonFromGM(int addr);
    Photon LoadPhotonFromGMGivenIndex(int index);

	class PhotonMap
	{
    	int pm_s_addr;
    	int num_pht;
    	const int pht_size;

    	bool sorted; // flag to identify whether photons are sorted
    	int _tid; // thread id
        
        int s_mat_addr;

        AABBox sbox; // scene bounding box, assume cornel box scene, all positive

        Photon LoadPhotonGivenID(int index);

    	bool BruteForceRadEst(HitRecord & hit, const float r, Color & res, float &omax_r2, int &nnei);
    	bool ZcurveRadEst(HitRecord &hit, const float r, Color &res);

    	int GetZcurveIndexGivenPhotonId(int index, Photon & pht);

	public:
    	PhotonMap();

		void Compute(const Scene & scene, const int threadId);
		void Sort(); // Sort photons according to z-curve
		bool RadianceEstimate(HitRecord & hit, const float r, Color & res, float &omax_r2, int &nnei);
        
        void Reset();
        
        inline void SetMaterialStartAddr(int iaddr);
        inline int PhotonAddr(int index) const;
	}; // class PhotonMap

	void PhotonSwap(Photon & pht1, Photon & pht2, int & zid1, int &zid2, const int addr1, const int addr2);
    
    // ********** Inline Functions *****************
    void PhotonMap::SetMaterialStartAddr(int addr)
    {
        s_mat_addr = addr;
    }

    int PhotonMap::PhotonAddr(int index) const
    {
    	return (pm_s_addr + 1 + index*pht_size);
    }

} // namespace AtomTrace

#endif // __ATOM_TRACE_PHOTON_MAP_H_
